import pandas as pd
from matplotlib import pyplot as plt
from sklearn.model_selection import train_test_split
from sklearn.ensemble import RandomForestRegressor
from sklearn.metrics import mean_squared_error
import seaborn as sns
from statsmodels.tsa.arima.model import ARIMA
import numpy as np

import warnings

from statsmodels.tsa.stattools import adfuller

warnings.filterwarnings('ignore')

"""
特征值（输入量）： 每届运动会参加的运动员数量
标签值（输出量）：

随机森林1.0 输出两个图：
1. 实际值与预测值的对比图
   - 作用：展示模型预测的金牌数量与真实金牌数量的匹配程度。如果点分布接近红色理想预测线（y=x），说明预测性能较好。
2. 预测误差的分布直方图
   - 作用：揭示模型预测误差的分布情况。如果误差集中在0附近，说明模型整体误差较小；误差的分布趋势和范围可以帮助判断模型的稳健性。
"""

MIN_GROWTH_RATE = 0.05
MIN_VAL = 1  #
medal_counts_df = pd.read_csv('data/summerOly_medal_counts_cleaned.csv', encoding='latin1')
athletes_df = pd.read_csv('data/summerOly_athletes.csv', encoding='latin1')

# 将 Year 转换成 int 型
medal_counts_df['Year'] = medal_counts_df['Year'].astype(int)
athletes_df['Year'] = athletes_df['Year'].astype(int)

# -------------------------------------------------------- #

# 每个国家（NOC）在某一年参加夏季奥运会的所有运动员的总人数。
athletes_count = athletes_df.groupby(['NOC', 'Year']).size().reset_index(name='Athletes')
merged_df = pd.merge(medal_counts_df, athletes_count, on=['NOC', 'Year'], how='left')

# 每个国家在每届奥运会中参与的不同运动项目的数量
events_count = athletes_df.groupby(['NOC', 'Year'])['Sport'].nunique().reset_index(name='Events')
merged_df = pd.merge(merged_df, events_count, on=['NOC', 'Year'], how='left')

# print(merged_df.head())

# -------------------------------------------------------- #

# 计算 2024年 每个国家 参加的运动员数量 和 参与的项目数
athletes_2024 = merged_df[merged_df['Year'] == 2024].groupby('NOC')['Athletes'].mean()
events_2024 = merged_df[merged_df['Year'] == 2024].groupby('NOC')['Events'].mean()

# 2024年 Mixed team（混合队）的数据被排除
athletes_2024 = athletes_2024[athletes_2024.index != 'Mixed team']
events_2024 = events_2024[events_2024.index != 'Mixed team']

# 对 athletes_2024 和 events_2024 中的缺失值进行填充。
# MIN_VAL 是之前定义的常量（通常为1），表示如果某个国家的运动员人数或独立项目数数据缺失，就用这个最小值填充。
athletes_2024.fillna(MIN_VAL, inplace=True)
events_2024.fillna(MIN_VAL, inplace=True)

# 处理异常值
athletes_2024[athletes_2024 < MIN_VAL] = MIN_VAL
events_2024[events_2024 < MIN_VAL] = MIN_VAL


# print(athletes_2024.head())
# print(events_2024.head())

# -------------------------------------------------------- #

# 对于每个国家/地区（NOC），计算其在某个列（如运动员人数 Athletes 或独立项目数 Events）上的年均增长率。

def calculate_growth_rate(df, column):
    # 分组后，如果只有一条记录或首条=0，返回MIN_GROWTH_RATE
    def _growth(group):
        if len(group) > 1 and group[column].iloc[0] != 0:
            ratio = (group[column].iloc[-1] / group[column].iloc[0]) ** (1 / (len(group) - 1)) - 1
            return ratio if ratio > 0 else MIN_GROWTH_RATE
        else:
            return MIN_GROWTH_RATE

    result = df.groupby('NOC').apply(_growth).reset_index(name=f'{column}_growth_rate')
    return result


# 计算增长率
athletes_growth = calculate_growth_rate(merged_df, 'Athletes')
events_growth = calculate_growth_rate(merged_df, 'Events')

athletes_2024_growth = athletes_growth.set_index('NOC')['Athletes_growth_rate']
events_2024_growth = events_growth.set_index('NOC')['Events_growth_rate']

# 通过计算年增长率来预测2028年的值
athletes_2028 = athletes_2024 * (1 + athletes_2024_growth)
events_2028 = events_2024 * (1 + events_2024_growth)

predictions_gold = {}
predictions_total = {}
predictions_interval_gold = {}
predictions_interval_total = {}

for country in merged_df['NOC'].unique():
    country_data = merged_df[merged_df['NOC'] == country]
    if len(country_data) < 2:
        print(f"Skipping {country}, insufficient data.")
        continue

    X = country_data[['Year', 'Athletes', 'Events']]
    y_gold = country_data['Gold']
    y_total = country_data['Total']
    X_train, X_test, y_train_gold, y_test_gold = train_test_split(X, y_gold, test_size=0.2, random_state=42)
    X_train_total, X_test_total, y_train_total, y_test_total = train_test_split(X, y_total, test_size=0.2,
                                                                                random_state=42)

    model_gold = RandomForestRegressor(n_estimators=300, random_state=42)
    model_total = RandomForestRegressor(n_estimators=300, random_state=42)
    model_gold.fit(X_train, y_train_gold)
    model_total.fit(X_train_total, y_train_total)
    mse_gold = mean_squared_error(y_test_gold, model_gold.predict(X_test))
    mse_total = mean_squared_error(y_test_total, model_total.predict(X_test_total))

    print(f"{country}: Gold MSE={mse_gold}, Total MSE={mse_total}")

    if country in athletes_2028.index and country in events_2028.index:
        X_2028 = pd.DataFrame({
            'Year': [2028],
            'Athletes': [athletes_2028[country]],
            'Events': [events_2028[country]]
        })
        X_2028 = X_2028.replace([np.inf, -np.inf], 0).fillna(0)

        pred_gold_2028 = model_gold.predict(X_2028)
        pred_total_2028 = model_total.predict(X_2028)

        # 不确定性区间
        interval_g = np.std([model_gold.predict(X_2028) for _ in range(100)])
        interval_t = np.std([model_total.predict(X_2028) for _ in range(100)])

        predictions_gold[country] = pred_gold_2028[0]
        predictions_total[country] = pred_total_2028[0]
        predictions_interval_gold[country] = interval_g
        predictions_interval_total[country] = interval_t


predictions_df = pd.DataFrame({
    'NOC': list(predictions_gold.keys()),
    'Predicted Gold Medals for 2028': list(predictions_gold.values()),
    'Predicted Total Medals for 2028': list(predictions_total.values()),
    'Gold Prediction Interval': list(predictions_interval_gold.values()),
    'Total Prediction Interval': list(predictions_interval_total.values())
})


# 按照预测的总奖牌数从高到低排序
predictions_df = predictions_df.sort_values(by='Predicted Total Medals for 2028', ascending=False)

# 保存排序后的结果到文件
predictions_df.to_csv('2028_predictions_by_country_sorted.csv', index=False)

print(predictions_df.head())  # 输出前几行查看排序结果

# 获取预测总奖牌数前四的国家
top_4_countries = predictions_df.head(4)['NOC']

# 初始化一个 2x2 的子图
fig, axes = plt.subplots(2, 2, figsize=(12, 10))
axes = axes.flatten()

for i, country in enumerate(top_4_countries):
    country_data = merged_df[merged_df['NOC'] == country]
    if len(country_data) < 2:
        print(f"Skipping {country}, insufficient data.")
        continue

    # 提取实际值和预测值
    actual_values = country_data['Total']
    years = country_data['Year']

    # 训练模型进行预测
    model = RandomForestRegressor(n_estimators=300, random_state=42)
    model.fit(country_data[['Year', 'Athletes', 'Events']], country_data['Total'])
    predicted_values = model.predict(country_data[['Year', 'Athletes', 'Events']])

    # 绘制散点图
    ax = axes[i]
    sns.scatterplot(x=years, y=actual_values, label='Actual', ax=ax, color='blue')
    sns.scatterplot(x=years, y=predicted_values, label='Predicted', ax=ax, color='red')

    # 标注国家名称
    ax.text(0.5, 0.9, f'{country}', fontsize=12, transform=ax.transAxes,
            ha='center', va='center', bbox=dict(boxstyle='round', facecolor='white', alpha=0.8))

    # 标题和图例
    ax.set_title(f'{country} - Actual vs Predicted Total Medals')
    ax.set_xlabel('Year')
    ax.set_ylabel('Medals')
    ax.legend()

# 调整布局
plt.tight_layout()
plt.show()